from functools import cached_property
import re
from typing import ClassVar
from typing_extensions import Self
import polars as pl
from pydantic import BaseModel, model_validator
from pysam import FastaFile  # pylint: disable=E0611
from .snv import SNV
from .transcript import Transcript
from .enum import Criterion, LoF, Strength
from .load_data import BEDFrame
from .splicing import Splicing


class PVS1(BaseModel):
    GenomeFasta: ClassVar[FastaFile]
    Domain: ClassVar[BEDFrame]
    Hotspot: ClassVar[BEDFrame]
    CuratedRegion: ClassVar[BEDFrame]
    ExonLofPopmax: ClassVar[BEDFrame]
    Pathogenic: ClassVar[pl.DataFrame]
    snv: SNV
    transcript: Transcript
    criterion: Criterion = Criterion.IC5
    strength_raw: Strength = Strength.Unmet

    class Config:
        arbitrary_types_allowed = True

    def __str__(self) -> str:
        return f"{self.criterion}|{self.strength_raw}|{self.strength}"

    @classmethod
    def default_str(cls) -> str:
        return f"{Criterion.NF0}|{Strength.Unmet}|{Strength.Unset}"

    @cached_property
    def is_nmd_target(self) -> bool:
        return self.snv.termination * 3 <= self.transcript.nmd_cutoff

    @cached_property
    def is_biologically_relevant(self) -> bool:
        return True

    @cached_property
    def is_critical_to_protein_func(self) -> bool:
        """
        PM1 mutational hot spot and/or critical and well-established functional domain.
        """
        if self.transcript.gene == "CDH1":
            return self.snv.termination <= 836
        start, end = (self.snv.pos - 1, self.transcript.cds_end) if self.transcript.strand == "+" else (self.transcript.cds_start, self.snv.pos)
        hotspot_row = self.Hotspot.fetch_row(self.snv.chrom, start, end)
        curated_region_row = self.CuratedRegion.fetch_row(self.snv.chrom, start, end)
        if not (curated_region_row or hotspot_row):
            domain_row = self.Domain.fetch_row(self.snv.chrom, start, end)
            if domain_row:
                missense_plp, missense_blb = domain_row["name"].split("|")[5:7]
                return (int(missense_blb) == 0 and int(missense_plp) >= 5) or (int(missense_blb) > 0 and int(missense_plp) / int(missense_blb) >= 10)
            return False
        return True

    @cached_property
    def exon_lofs_are_frequent_in_pop(self) -> bool:
        """
        NM_153818.1.4|1-2338205-G-A:1.19e-04|1-2338230-C-CT:5.62e-05
        :return:
        """
        exon_lof_row = self.ExonLofPopmax.fetch_row(self.snv.chrom, self.snv.pos - 5, self.snv.pos + 5)
        if exon_lof_row:
            max_lof = exon_lof_row["name"].split("|")[1]
            return float(max_lof.split(":")[1]) > 0.001
        return False

    @cached_property
    def lof_removes_more_than_10_percent_of_protein(self) -> bool:
        """
        LoF variant removes >10% of protein.
        """
        m = re.compile(r"p\.\D+(\d+)(\D+fs)?(\*|X|Ter)(\d+)?").search(self.snv.aa_change)
        aa_pos = int(m.group(1)) if m else -1
        protein_size = self.transcript.cds_size / 3
        return aa_pos > 0 and (protein_size - aa_pos) / protein_size > 0.1

    @cached_property
    def pvs1_start_codon(self) -> tuple[Strength, Criterion]:
        records = self.transcript.closest_potential_start_codon
        if records.is_empty():
            return Strength.VeryStrong, Criterion.IC0
        records = records.join(self.Pathogenic, on=["chrom", "pos"], how="left").fill_null(0)
        variant_score = records["count"].sum()
        if variant_score < 1:
            return Strength.Supporting, Criterion.IC4
        if 1 <= variant_score <= 3:
            return Strength.Moderate, Criterion.IC3
        if 3 < variant_score <= 6:
            return Strength.Strong, Criterion.IC2
        return Strength.VeryStrong, Criterion.IC1

    @cached_property
    def strength(self) -> Strength:
        if self.transcript.gene == "MYH7":
            return Strength.Moderate if self.strength_raw.value >= 2 else self.strength_raw
        match self.transcript.level:
            case "L0":
                strength = self.strength_raw
            case "L1":
                strength = self.strength_raw.downgrade(1)
            case "L2":
                strength = self.strength_raw.downgrade(2)
            case "L3":
                strength = Strength.Unmet
            case _:
                strength = Strength.Unset
        return strength

    @model_validator(mode="after")
    def verify_pvs1(self) -> Self:
        if self.transcript.cds_size == 0:
            self.criterion, self.strength_raw = Criterion.NF0, Strength.Unmet
            return self
        if self.snv.lof in [LoF.Nonsense, LoF.FrameShift]:
            if self.transcript.gene == "PTEN" and self.snv.termination < 374:
                self.criterion, self.strength_raw = Criterion.PTEN, Strength.VeryStrong
            else:
                if self.is_nmd_target:
                    if self.is_biologically_relevant:
                        self.criterion, self.strength_raw = (
                            Criterion.NF1,
                            Strength.VeryStrong,
                        )
                    else:
                        self.criterion, self.strength_raw = (
                            Criterion.NF2,
                            Strength.Unmet,
                        )
                else:
                    if self.is_critical_to_protein_func:
                        self.criterion, self.strength_raw = (
                            Criterion.NF3,
                            Strength.Strong,
                        )
                    else:
                        if self.exon_lofs_are_frequent_in_pop or not self.is_biologically_relevant:
                            self.criterion, self.strength_raw = (
                                Criterion.NF4,
                                Strength.Unmet,
                            )
                        elif self.lof_removes_more_than_10_percent_of_protein:
                            self.criterion, self.strength_raw = (
                                Criterion.NF5,
                                Strength.Strong,
                            )
                        else:
                            self.criterion, self.strength_raw = (
                                Criterion.NF6,
                                Strength.Moderate,
                            )
        elif self.snv.lof == LoF.Splicing:
            if self.transcript.gene == "PTEN":
                m = re.search(r"c\.(\d+)", self.snv.aa_change)
                if m and int(m.group(1)) / 3 < 374:
                    self.criterion, self.strength_raw = (
                        Criterion.PTEN,
                        Strength.VeryStrong,
                    )
                    return self
            # CDH1 Guideline: Use the strong strength of evidence for canonical splice sites.
            if self.transcript.gene == "CDH1":
                self.criterion, self.strength_raw = Criterion.CDH1, Strength.Strong
            else:
                splice = Splicing(snv=self.snv, transcript=self.transcript)
                if splice.preserves_reading_frame:
                    if splice.is_critical_to_protein_func:
                        self.criterion, self.strength_raw = (
                            Criterion.SS10,
                            Strength.Strong,
                        )
                    elif self.exon_lofs_are_frequent_in_pop or not self.is_biologically_relevant:
                        self.criterion, self.strength_raw = (
                            Criterion.SS7,
                            Strength.Unmet,
                        )
                    elif splice.variant_removes_10_percent_of_protein:
                        self.criterion, self.strength_raw = (
                            Criterion.SS8,
                            Strength.Strong,
                        )
                    else:
                        self.criterion, self.strength_raw = (
                            Criterion.SS9,
                            Strength.Moderate,
                        )
                elif splice.is_undergo_nmd:
                    if self.is_biologically_relevant:
                        self.criterion, self.strength_raw = (
                            Criterion.SS1,
                            Strength.VeryStrong,
                        )
                    else:
                        self.criterion, self.strength_raw = (
                            Criterion.SS2,
                            Strength.Unmet,
                        )
                else:
                    if splice.is_critical_to_protein_func:
                        self.criterion, self.strength_raw = (
                            Criterion.SS3,
                            Strength.Strong,
                        )
                    elif self.exon_lofs_are_frequent_in_pop or not self.is_biologically_relevant:
                        self.criterion, self.strength_raw = (
                            Criterion.SS4,
                            Strength.Unmet,
                        )
                    elif splice.variant_removes_10_percent_of_protein:
                        self.criterion, self.strength_raw = (
                            Criterion.SS5,
                            Strength.Strong,
                        )
                    else:
                        self.criterion, self.strength_raw = (
                            Criterion.SS6,
                            Strength.Moderate,
                        )
        elif self.snv.lof == LoF.StartLoss:
            self.strength_raw, self.criterion = self.pvs1_start_codon
        else:
            self.criterion, self.strength_raw = Criterion.IC5, Strength.Unmet
        return self
